// 用类描述实现前缀树。
// 测试链接 : https://leetcode.cn/problems/implement-trie-ii-prefix-tree/
// https://www.nowcoder.com/practice/7f8a8553ddbf4eaab749ec988726702b
public class TrieTree1 {
    class TrieNode{
        public int pass;
        public int end;
        public TrieNode[] nexts;

        public TrieNode(){
            pass = 0;
            end = 0;
            nexts = new TrieNode[26];
        }
    }

    private TrieNode root;

    public TrieTree1() {
        root = new TrieNode();
    }

    public void insert(String word) {
        TrieNode cur = root;
        cur.pass++;

        for(int i = 0, path; i < word.length(); i++){
            path = word.charAt(i) - 'a';

            if(cur.nexts[path] == null){
                cur.nexts[path] = new TrieNode();
            }
            cur = cur.nexts[path];
            cur.pass++;
        }
        cur.end++;
    }

    // 如果之前word插入过前缀树，那么此时删掉一次
    // 如果之前word没有插入过前缀树，那么什么也不做
    public void erase(String word) {
        if(countWordsEqualTo(word) == 0){  //不存在当前这个字符串直接结束方法
            return;
        }

        TrieNode cur = root;
        cur.pass--;

        for(int i = 0, path; i < word.length(); i++){
            path = word.charAt(i) - 'a';

            if(--cur.nexts[path].pass == 0){
                cur.nexts[path] = null;
                return;
            }
            cur = cur.nexts[path];
        }
        cur.end--;
    }

    // 查询前缀树里，word单词出现了几次
    public int countWordsEqualTo(String word) {
        TrieNode cur = root;

        for(int i = 0, path; i < word.length(); i++){
            path = word.charAt(i) - 'a';

            if(cur.nexts[path] == null){
                return 0;
            }
            cur = cur.nexts[path];
        }
        return cur.end;
    }

    // 查询前缀树里，有多少单词以pre做前缀
    public int countWordsStartingWith(String pre) {
        TrieNode cur = root;;

        for(int i = 0, path; i < pre.length(); i++){
            path = pre.charAt(i) - 'a';
            if(cur.nexts[path] == null){
                return 0;
            }
            cur = cur.nexts[path];
        }
        return cur.pass;
    }
}
